#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-
import datetime
import time

import backtrader as bt
from backtrader.feeds import InfluxDB
from influxdb.client import InfluxDBClient as idbclient
from backtrader.utils import date2num
import datetime as dt
import utils

from influxdb.exceptions import InfluxDBClientError

TIMEFRAMES = dict(
    (
        (bt.TimeFrame.Seconds, 's'),
        (bt.TimeFrame.Minutes, 'm'),
        (bt.TimeFrame.Days, 'd'),
        (bt.TimeFrame.Weeks, 'w'),
        (bt.TimeFrame.Months, 'm'),
        (bt.TimeFrame.Years, 'y'),
    )
)


class InfluxDB(InfluxDB):
    frompackages = (
        ('influxdb', [('InfluxDBClient', 'idbclient')]),
        ('influxdb.exceptions', 'InfluxDBClientError')
    )

    params = (
        ('host', '101.36.109.137'),
        ('port', '8086'),
        ('username', 'chrise'),
        ('password', '.cf@135589'),
        ('database', None),
        ('timeframe', bt.TimeFrame.Days),
        ('startdate', None),
        ('high', 'high_p'),
        ('low', 'low_p'),
        ('open', 'open_p'),
        ('close', 'close_p'),
        ('volume', 'volume'),
        ('ointerest', 'oi'),
        ('drop_newest', True),
        ('historical', False)
    )

    # States for the Finite State Machine in _load
    _ST_FROM, _ST_START, _ST_LIVE, _ST_HISTORBACK, _ST_OVER = range(5)

    def __init__(self):
        self._state = self._ST_LIVE
        self._lasttime = ''

    def start(self):
        super(InfluxDB, self).start()
        try:
            self.ndb = idbclient(self.p.host, self.p.port, self.p.username,
                                 self.p.password, self.p.database)
        except InfluxDBClientError as err:
            print('Failed to establish connection to InfluxDB: %s' % err)

        tf = '{multiple}{timeframe}'.format(
            multiple=(self.p.compression if self.p.compression else 1),
            timeframe=TIMEFRAMES.get(self.p.timeframe, 'd'))

        if not self.p.startdate:
            st = '<= now()'
        else:
            st = '>= \'%s\'' % self.p.startdate

        # The query could already consider parameters like fromdate and todate
        # to have the database skip them and not the internal code
        qstr = ('SELECT mean("{open_f}") AS "open", mean("{high_f}") AS "high", '
                'mean("{low_f}") AS "low", mean("{close_f}") AS "close", '
                'mean("{vol_f}") AS "volume", mean("{oi_f}") AS "openinterest" '
                'FROM "{dataname}" '
                'WHERE time {begin} '
                'GROUP BY time({timeframe}) fill(none)').format(
            open_f=self.p.open, high_f=self.p.high,
            low_f=self.p.low, close_f=self.p.close,
            vol_f=self.p.volume, oi_f=self.p.ointerest,
            timeframe=tf, begin=st, dataname=self.p.dataname)

        try:
            dbars = list(self.ndb.query(qstr).get_points())
        except InfluxDBClientError as err:
            print('InfluxDB query failed: %s' % err)

        self.biter = iter(self._formatdbars(dbars))

    def _formatdbars(self, dbars):
        if (self.p.drop_newest):
            dbars = dbars[0:-1]
        return dbars

    def _load(self):
        _retry_times = 1
        _sleep_seconds = 30
        while True:
            bar = next(self.biter, {})

            if bar:
                if _retry_times > 1: print(bar)
                break

            # 尝试30次未获得新数据，返回False
            if _retry_times > 3600 / _sleep_seconds:
                print('Retry {} times,return false.'.format(_retry_times))
                return False

            time.sleep(_sleep_seconds)
            print(
                '{} Load bar at time {},retry {} times'.format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                                                               utils.utc2local(self._lasttime,
                                                                               UTC_FORMAT="%Y-%m-%dT%H:%M:%SZ"),
                                                               _retry_times
                                                               ))
            _retry_times = _retry_times + 1
            self.getbars()

        self.l.datetime[0] = date2num(dt.datetime.strptime(bar['time'],
                                                           '%Y-%m-%dT%H:%M:%SZ'))

        self.l.open[0] = bar['open']
        self.l.high[0] = bar['high']
        self.l.low[0] = bar['low']
        self.l.close[0] = bar['close']
        self.l.volume[0] = bar['volume']

        self._lasttime = bar['time']

        return True

    def getbars(self):
        tf = '{multiple}{timeframe}'.format(
            multiple=(self.p.compression if self.p.compression else 1),
            timeframe=TIMEFRAMES.get(self.p.timeframe, 'd'))

        st = '> \'%s\'' % self._lasttime

        # The query could already consider parameters like fromdate and todate
        # to have the database skip them and not the internal code
        qstr = ('SELECT mean("{open_f}") AS "open", mean("{high_f}") AS "high", '
                'mean("{low_f}") AS "low", mean("{close_f}") AS "close", '
                'mean("{vol_f}") AS "volume", mean("{oi_f}") AS "openinterest" '
                'FROM "{dataname}" '
                'WHERE time {begin} '
                'GROUP BY time({timeframe}) fill(none)').format(
            open_f=self.p.open, high_f=self.p.high,
            low_f=self.p.low, close_f=self.p.close,
            vol_f=self.p.volume, oi_f=self.p.ointerest,
            timeframe=tf, begin=st, dataname=self.p.dataname)

        try:
            dbars = list(self.ndb.query(qstr).get_points())
        except InfluxDBClientError as err:
            print('InfluxDB query failed: %s' % err)

        self.biter = iter(self._formatdbars(dbars))

    def haslivedata(self):
        return self._state == self._ST_LIVE

    def islive(self):
        return not self.p.historical
